#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/point_cloud2.hpp>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/voxel_grid.h>
#include <OpenCL/opencl.h>
#include <cmath>
#include <fstream>
#include <sstream>
#include <vector>
#include <memory>
#include <mutex>

using namespace std;

// OpenCL相关全局变量 - 共享上下文避免冲突
cl_platform_id platform;
cl_device_id device;
cl_context context;
cl_command_queue queue;
cl_program program;
cl_kernel transform_kernel;    // 坐标变换内核
cl_kernel ray_filter_kernel;   // 光线滤波内核
std::mutex opencl_mutex;       // 互斥锁防止并发访问冲突

// 光线滤波参数
struct RayFilterParams {
    float sensor_height = 1.8f;       // 传感器高度(m)
    float max_distance = 50.0f;       // 最大处理距离(m)
    float min_distance = 1.0f;        // 最小处理距离(m)
    float angle_resolution = 0.01745f;// 角度分辨率(rad) ≈1°
    float max_height_diff = 0.15f;    // 最大高度差阈值(m)
    float initial_height_threshold = 0.5f; // 初始地面点高度阈值(m)
};

class PointCloudProcessorNode : public rclcpp::Node {
public:
    PointCloudProcessorNode() : Node("pointcloud_processor_node") {
        // 声明坐标变换参数
        declare_parameter<float>("rotation_angle", 0.0f); // 旋转角度(度)
        declare_parameter<float>("translation_x", 0.0f);
        declare_parameter<float>("translation_y", 0.0f);
        declare_parameter<float>("translation_z", 0.0f);
        
        // 声明光线滤波参数
        declare_parameter<float>("sensor_height", 1.8f);
        declare_parameter<float>("max_distance", 50.0f);
        declare_parameter<float>("min_distance", 1.0f);
        declare_parameter<float>("angle_resolution", 0.01745f);
        declare_parameter<float>("max_height_diff", 0.15f);
        declare_parameter<float>("initial_height_threshold", 0.5f);
        
        // 创建订阅者
        transform_sub_ = create_subscription<sensor_msgs::msg::PointCloud2>(
            "/points_to_transform", 10,
            bind(&PointCloudProcessorNode::transform_callback, this, placeholders::_1));
        
        filter_sub_ = create_subscription<sensor_msgs::msg::PointCloud2>(
            "/points_to_filter", 10,
            bind(&PointCloudProcessorNode::filter_callback, this, placeholders::_1));
        
        // 创建发布者
        transformed_pub_ = create_publisher<sensor_msgs::msg::PointCloud2>(
            "/transformed_points", 10);
        
        ground_pub_ = create_publisher<sensor_msgs::msg::PointCloud2>(
            "/ground_points", 10);
        
        non_ground_pub_ = create_publisher<sensor_msgs::msg::PointCloud2>(
            "/non_ground_points", 10);
        
        // 初始化OpenCL - 只初始化一次，两个内核共享上下文
        if (!init_opencl()) {
            RCLCPP_ERROR(get_logger(), "Failed to initialize OpenCL");
            rclcpp::shutdown();
        }
        
        RCLCPP_INFO(get_logger(), "PointCloud Processor Node initialized with OpenCL acceleration");
    }
    
    ~PointCloudProcessorNode() {
        // 释放OpenCL资源
        clReleaseKernel(transform_kernel);
        clReleaseKernel(ray_filter_kernel);
        clReleaseProgram(program);
        clReleaseCommandQueue(queue);
        clReleaseContext(context);
        clReleaseDevice(device);
    }

private:
    rclcpp::Subscription<sensor_msgs::msg::PointCloud2>::SharedPtr transform_sub_;
    rclcpp::Subscription<sensor_msgs::msg::PointCloud2>::SharedPtr filter_sub_;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr transformed_pub_;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr ground_pub_;
    rclcpp::Publisher<sensor_msgs::msg::PointCloud2>::SharedPtr non_ground_pub_;
    RayFilterParams filter_params_;
    
    // 初始化OpenCL环境 - 创建两个内核并共享上下文
    bool init_opencl() {
        cl_int err;
        
        // 获取平台
        err = clGetPlatformIDs(1, &platform, NULL);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to get OpenCL platform");
            return false;
        }
        
        // 获取设备
        err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, NULL);
        if (err != CL_SUCCESS) {
            RCLCPP_WARN(get_logger(), "No GPU found, trying CPU");
            err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_CPU, 1, &device, NULL);
            if (err != CL_SUCCESS) {
                RCLCPP_ERROR(get_logger(), "Failed to get OpenCL device");
                return false;
            }
        }
        
        // 创建上下文 - 两个内核共享
        context = clCreateContext(NULL, 1, &device, NULL, NULL, &err);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to create OpenCL context");
            return false;
        }
        
        // 创建命令队列 - 两个内核共享
        queue = clCreateCommandQueue(context, device, 0, &err);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to create OpenCL command queue");
            return false;
        }
        
        // 加载并编译包含两个内核的程序
        const string kernel_source = R"(
            // 坐标变换内核 - 绕Z轴旋转和平移
            __kernel void point_transform(
                __global const float* in_x,
                __global const float* in_y,
                __global const float* in_z,
                __global float* out_x,
                __global float* out_y,
                __global float* out_z,
                const int num_points,
                const float cos_theta,
                const float sin_theta,
                const float tx,
                const float ty,
                const float tz
            ) {
                int i = get_global_id(0);
                if (i >= num_points) return;
                
                // 读取输入点
                float x = in_x[i];
                float y = in_y[i];
                float z = in_z[i];
                
                // 旋转变换 (绕Z轴)
                float x_rot = x * cos_theta - y * sin_theta;
                float y_rot = x * sin_theta + y * cos_theta;
                
                // 平移变换
                out_x[i] = x_rot + tx;
                out_y[i] = y_rot + ty;
                out_z[i] = z + tz;
            }

            // 光线法滤波内核
            __kernel void ray_ground_filter(
                __global const float* x,
                __global const float* y,
                __global const float* z,
                __global uchar* is_ground,
                const int num_points,
                const float sensor_height,
                const float max_distance,
                const float min_distance,
                const float angle_resolution,
                const float max_height_diff,
                const float initial_height_threshold
            ) {
                int i = get_global_id(0);
                if (i >= num_points) return;
                
                // 计算点的距离和角度
                float x_val = x[i];
                float y_val = y[i];
                float z_val = z[i];
                
                float dist = sqrt(x_val*x_val + y_val*y_val);
                
                // 过滤超出距离范围的点
                if (dist < min_distance || dist > max_distance) {
                    is_ground[i] = 0;
                    return;
                }
                
                // 计算水平角度 (弧度)
                float angle = atan2(y_val, x_val);
                if (angle < 0) angle += 2.0f * M_PI_F;
                
                // 计算所属的射线索引
                int ray_index = (int)(angle / angle_resolution);
                
                // 存储每条射线上的点索引
                __local int* ray_counts;
                if (get_local_id(0) == 0) {
                    int num_rays = (int)(2.0f * M_PI_F / angle_resolution) + 1;
                    ray_counts = (int*)malloc(num_rays * sizeof(int));
                    for (int j = 0; j < num_rays; j++) ray_counts[j] = 0;
                }
                barrier(CLK_LOCAL_MEM_FENCE);
                
                // 统计每条射线的点数量
                atomic_inc(&ray_counts[ray_index]);
                barrier(CLK_LOCAL_MEM_FENCE);
                
                // 为每条射线分配内存并存储点索引和距离
                __local int* ray_points;
                __local float* ray_distances;
                if (get_local_id(0) == 0) {
                    int total = 0;
                    for (int j = 0; j < (int)(2.0f * M_PI_F / angle_resolution) + 1; j++) {
                        total += ray_counts[j];
                    }
                    ray_points = (int*)malloc(total * sizeof(int));
                    ray_distances = (float*)malloc(total * sizeof(float));
                    
                    // 重置计数以重新填充
                    int num_rays = (int)(2.0f * M_PI_F / angle_resolution) + 1;
                    for (int j = 0; j < num_rays; j++) ray_counts[j] = 0;
                }
                barrier(CLK_LOCAL_MEM_FENCE);
                
                // 填充每条射线的点数据
                int pos = atomic_inc(&ray_counts[ray_index]);
                int offset = 0;
                for (int j = 0; j < ray_index; j++) {
                    offset += ray_counts[j];
                }
                ray_points[offset + pos] = i;
                ray_distances[offset + pos] = dist;
                barrier(CLK_LOCAL_MEM_FENCE);
                
                // 每条射线由一个工作项处理排序和地面判断
                if (get_local_id(0) == 0) {
                    int num_rays = (int)(2.0f * M_PI_F / angle_resolution) + 1;
                    for (int ray = 0; ray < num_rays; ray++) {
                        int count = ray_counts[ray];
                        if (count == 0) continue;
                        
                        // 计算该射线在数组中的起始位置
                        int start = 0;
                        for (int j = 0; j < ray; j++) {
                            start += ray_counts[j];
                        }
                        
                        // 按距离排序
                        for (int a = 0; a < count; a++) {
                            for (int b = a + 1; b < count; b++) {
                                if (ray_distances[start + a] > ray_distances[start + b]) {
                                    // 交换距离
                                    float temp_dist = ray_distances[start + a];
                                    ray_distances[start + a] = ray_distances[start + b];
                                    ray_distances[start + b] = temp_dist;
                                    
                                    // 交换点索引
                                    int temp_idx = ray_points[start + a];
                                    ray_points[start + a] = ray_points[start + b];
                                    ray_points[start + b] = temp_idx;
                                }
                            }
                        }
                        
                        // 识别地面点
                        float prev_z = -sensor_height;
                        bool found_initial = false;
                        
                        for (int j = 0; j < count; j++) {
                            int point_idx = ray_points[start + j];
                            float current_z = z[point_idx];
                            float current_dist = ray_distances[start + j];
                            
                            // 寻找初始地面点
                            if (!found_initial) {
                                if (fabs(current_z + sensor_height) < initial_height_threshold) {
                                    is_ground[point_idx] = 1;
                                    prev_z = current_z;
                                    found_initial = true;
                                } else {
                                    is_ground[point_idx] = 0;
                                }
                            } else {
                                // 基于前一个地面点判断当前点
                                if (current_z > prev_z - max_height_diff) {
                                    is_ground[point_idx] = 1;
                                    prev_z = current_z;
                                } else {
                                    is_ground[point_idx] = 0;
                                }
                            }
                        }
                    }
                    
                    // 释放本地内存
                    free(ray_counts);
                    free(ray_points);
                    free(ray_distances);
                }
            }
        )";
        
        // 创建程序
        const char* source = kernel_source.c_str();
        program = clCreateProgramWithSource(context, 1, &source, NULL, &err);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to create OpenCL program");
            return false;
        }
        
        // 编译程序
        err = clBuildProgram(program, 1, &device, "-cl-mad-enable", NULL, NULL);
        if (err != CL_SUCCESS) {
            // 获取编译错误信息
            size_t log_size;
            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &log_size);
            char* log = (char*)malloc(log_size);
            clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, log, NULL);
            RCLCPP_ERROR(get_logger(), "OpenCL program build error:\n%s", log);
            free(log);
            return false;
        }
        
        // 创建两个内核
        transform_kernel = clCreateKernel(program, "point_transform", &err);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to create transform kernel");
            return false;
        }
        
        ray_filter_kernel = clCreateKernel(program, "ray_ground_filter", &err);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to create ray filter kernel");
            return false;
        }
        
        RCLCPP_INFO(get_logger(), "OpenCL initialized with both kernels successfully");
        return true;
    }
    
    // 坐标变换回调函数 - 使用OpenCL加速
    void transform_callback(const sensor_msgs::msg::PointCloud2::SharedPtr msg) {
        // 获取参数
        float rotation_angle = get_parameter("rotation_angle").as_double();
        float tx = get_parameter("translation_x").as_double();
        float ty = get_parameter("translation_y").as_double();
        float tz = get_parameter("translation_z").as_double();
        
        // 转换为PCL点云
        pcl::PointCloud<pcl::PointXYZ> cloud;
        pcl::fromROSMsg(*msg, cloud);
        
        if (cloud.empty()) {
            RCLCPP_WARN(get_logger(), "Received empty point cloud for transformation");
            return;
        }
        
        // 准备数据用于OpenCL处理
        int num_points = cloud.size();
        vector<float> in_x(num_points);
        vector<float> in_y(num_points);
        vector<float> in_z(num_points);
        vector<float> out_x(num_points);
        vector<float> out_y(num_points);
        vector<float> out_z(num_points);
        
        for (int i = 0; i < num_points; i++) {
            in_x[i] = cloud[i].x;
            in_y[i] = cloud[i].y;
            in_z[i] = cloud[i].z;
        }
        
        // 计算旋转角度的正弦和余弦值
        float rad = rotation_angle * M_PI / 180.0f;
        float cos_theta = cos(rad);
        float sin_theta = sin(rad);
        
        // 使用互斥锁确保OpenCL资源不会被同时访问
        std::lock_guard<std::mutex> lock(opencl_mutex);
        
        // 创建OpenCL缓冲区
        cl_int err;
        cl_mem in_x_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                        sizeof(float) * num_points, in_x.data(), &err);
        cl_mem in_y_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                        sizeof(float) * num_points, in_y.data(), &err);
        cl_mem in_z_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                        sizeof(float) * num_points, in_z.data(), &err);
        cl_mem out_x_buf = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                         sizeof(float) * num_points, NULL, &err);
        cl_mem out_y_buf = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                         sizeof(float) * num_points, NULL, &err);
        cl_mem out_z_buf = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                         sizeof(float) * num_points, NULL, &err);
        
        // 设置变换内核参数
        err = clSetKernelArg(transform_kernel, 0, sizeof(cl_mem), &in_x_buf);
        err |= clSetKernelArg(transform_kernel, 1, sizeof(cl_mem), &in_y_buf);
        err |= clSetKernelArg(transform_kernel, 2, sizeof(cl_mem), &in_z_buf);
        err |= clSetKernelArg(transform_kernel, 3, sizeof(cl_mem), &out_x_buf);
        err |= clSetKernelArg(transform_kernel, 4, sizeof(cl_mem), &out_y_buf);
        err |= clSetKernelArg(transform_kernel, 5, sizeof(cl_mem), &out_z_buf);
        err |= clSetKernelArg(transform_kernel, 6, sizeof(int), &num_points);
        err |= clSetKernelArg(transform_kernel, 7, sizeof(float), &cos_theta);
        err |= clSetKernelArg(transform_kernel, 8, sizeof(float), &sin_theta);
        err |= clSetKernelArg(transform_kernel, 9, sizeof(float), &tx);
        err |= clSetKernelArg(transform_kernel, 10, sizeof(float), &ty);
        err |= clSetKernelArg(transform_kernel, 11, sizeof(float), &tz);
        
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to set transform kernel arguments");
            return;
        }
        
        // 执行变换内核
        size_t global_work_size = num_points;
        size_t local_work_size = 256;
        if (local_work_size > global_work_size) local_work_size = global_work_size;
        
        err = clEnqueueNDRangeKernel(queue, transform_kernel, 1, NULL,
                                    &global_work_size, &local_work_size,
                                    0, NULL, NULL);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to execute transform kernel");
            return;
        }
        
        // 读取变换结果
        err = clEnqueueReadBuffer(queue, out_x_buf, CL_TRUE, 0,
                                 sizeof(float) * num_points, out_x.data(), 0, NULL, NULL);
        err |= clEnqueueReadBuffer(queue, out_y_buf, CL_TRUE, 0,
                                 sizeof(float) * num_points, out_y.data(), 0, NULL, NULL);
        err |= clEnqueueReadBuffer(queue, out_z_buf, CL_TRUE, 0,
                                 sizeof(float) * num_points, out_z.data(), 0, NULL, NULL);
        
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to read transform results");
            return;
        }
        
        // 释放缓冲区
        clReleaseMemObject(in_x_buf);
        clReleaseMemObject(in_y_buf);
        clReleaseMemObject(in_z_buf);
        clReleaseMemObject(out_x_buf);
        clReleaseMemObject(out_y_buf);
        clReleaseMemObject(out_z_buf);
        
        // 构建变换后的点云
        pcl::PointCloud<pcl::PointXYZ> transformed_cloud;
        transformed_cloud.header = cloud.header;
        
        for (int i = 0; i < num_points; i++) {
            pcl::PointXYZ point;
            point.x = out_x[i];
            point.y = out_y[i];
            point.z = out_z[i];
            transformed_cloud.push_back(point);
        }
        
        // 发布变换后的点云
        sensor_msgs::msg::PointCloud2 output_msg;
        pcl::toROSMsg(transformed_cloud, output_msg);
        transformed_pub_->publish(output_msg);
    }
    
    // 滤波回调函数 - 使用OpenCL加速
    void filter_callback(const sensor_msgs::msg::PointCloud2::SharedPtr msg) {
        // 更新滤波参数
        filter_params_.sensor_height = get_parameter("sensor_height").as_double();
        filter_params_.max_distance = get_parameter("max_distance").as_double();
        filter_params_.min_distance = get_parameter("min_distance").as_double();
        filter_params_.angle_resolution = get_parameter("angle_resolution").as_double();
        filter_params_.max_height_diff = get_parameter("max_height_diff").as_double();
        filter_params_.initial_height_threshold = get_parameter("initial_height_threshold").as_double();
        
        // 转换为PCL点云
        pcl::PointCloud<pcl::PointXYZ> cloud;
        pcl::fromROSMsg(*msg, cloud);
        
        if (cloud.empty()) {
            RCLCPP_WARN(get_logger(), "Received empty point cloud for filtering");
            return;
        }
        
        // 准备数据用于OpenCL处理
        int num_points = cloud.size();
        vector<float> x(num_points);
        vector<float> y(num_points);
        vector<float> z(num_points);
        
        for (int i = 0; i < num_points; i++) {
            x[i] = cloud[i].x;
            y[i] = cloud[i].y;
            z[i] = cloud[i].z;
        }
        
        vector<unsigned char> is_ground(num_points, 0);
        
        // 使用互斥锁确保OpenCL资源不会被同时访问
        std::lock_guard<std::mutex> lock(opencl_mutex);
        
        // 创建OpenCL缓冲区
        cl_int err;
        cl_mem x_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                     sizeof(float) * num_points, x.data(), &err);
        cl_mem y_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                     sizeof(float) * num_points, y.data(), &err);
        cl_mem z_buf = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                                     sizeof(float) * num_points, z.data(), &err);
        cl_mem ground_buf = clCreateBuffer(context, CL_MEM_WRITE_ONLY,
                                          sizeof(unsigned char) * num_points, NULL, &err);
        
        // 设置滤波内核参数
        err = clSetKernelArg(ray_filter_kernel, 0, sizeof(cl_mem), &x_buf);
        err |= clSetKernelArg(ray_filter_kernel, 1, sizeof(cl_mem), &y_buf);
        err |= clSetKernelArg(ray_filter_kernel, 2, sizeof(cl_mem), &z_buf);
        err |= clSetKernelArg(ray_filter_kernel, 3, sizeof(cl_mem), &ground_buf);
        err |= clSetKernelArg(ray_filter_kernel, 4, sizeof(int), &num_points);
        err |= clSetKernelArg(ray_filter_kernel, 5, sizeof(float), &filter_params_.sensor_height);
        err |= clSetKernelArg(ray_filter_kernel, 6, sizeof(float), &filter_params_.max_distance);
        err |= clSetKernelArg(ray_filter_kernel, 7, sizeof(float), &filter_params_.min_distance);
        err |= clSetKernelArg(ray_filter_kernel, 8, sizeof(float), &filter_params_.angle_resolution);
        err |= clSetKernelArg(ray_filter_kernel, 9, sizeof(float), &filter_params_.max_height_diff);
        err |= clSetKernelArg(ray_filter_kernel, 10, sizeof(float), &filter_params_.initial_height_threshold);
        
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to set filter kernel arguments");
            return;
        }
        
        // 执行滤波内核
        size_t global_work_size = num_points;
        size_t local_work_size = 256;
        if (local_work_size > global_work_size) local_work_size = global_work_size;
        
        err = clEnqueueNDRangeKernel(queue, ray_filter_kernel, 1, NULL,
                                    &global_work_size, &local_work_size,
                                    0, NULL, NULL);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to execute filter kernel");
            return;
        }
        
        // 读取滤波结果
        err = clEnqueueReadBuffer(queue, ground_buf, CL_TRUE, 0,
                                 sizeof(unsigned char) * num_points,
                                 is_ground.data(), 0, NULL, NULL);
        if (err != CL_SUCCESS) {
            RCLCPP_ERROR(get_logger(), "Failed to read filter results");
            return;
        }
        
        // 释放缓冲区
        clReleaseMemObject(x_buf);
        clReleaseMemObject(y_buf);
        clReleaseMemObject(z_buf);
        clReleaseMemObject(ground_buf);
        
        // 分离地面点和非地面点
        pcl::PointCloud<pcl::PointXYZ> ground_cloud, non_ground_cloud;
        ground_cloud.header = cloud.header;
        non_ground_cloud.header = cloud.header;
        
        for (int i = 0; i < num_points; i++) {
            if (is_ground[i]) {
                ground_cloud.push_back(cloud[i]);
            } else {
                non_ground_cloud.push_back(cloud[i]);
            }
        }
        
        RCLCPP_INFO(get_logger(), "Filtered: %d ground points, %d non-ground points",
                   ground_cloud.size(), non_ground_cloud.size());
        
        // 发布结果
        sensor_msgs::msg::PointCloud2 ground_msg, non_ground_msg;
        pcl::toROSMsg(ground_cloud, ground_msg);
        pcl::toROSMsg(non_ground_cloud, non_ground_msg);
        
        ground_pub_->publish(ground_msg);
        non_ground_pub_->publish(non_ground_msg);
    }
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);
    rclcpp::spin(make_shared<PointCloudProcessorNode>());
    rclcpp::shutdown();
    return 0;
}
